# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .base import BaseClassificationHead
from utils import CLASSIFICATION_HEADS
from sklearn.linear_model import LogisticRegression
import numpy as np


@CLASSIFICATION_HEADS.register_module("lr")
class LR(BaseClassificationHead):
    def __init__(self, label_name, head_params=None):
        super().__init__(label_name, head_params)
        self.model = LogisticRegression(**self.head_params)

    def train(self, datasets):
        self.logger.info(f"start train {self.model_name}")
        assert "positive" in datasets and "negative" in datasets, "datasets must contain positive and negative"
        positive_dataset = datasets["positive"]
        negative_dataset = datasets["negative"]
        # 正负例数据集格式：item_id: embedding
        positive_id_set = set(positive_dataset.keys())
        negative_id_set = set(negative_dataset.keys())
        # 正例和负例的交集，如果有交集，则说明有重复的item_id
        intersection_id_set = positive_id_set & negative_id_set
        if len(intersection_id_set) > 0:
            self.logger.warning(f"found {len(intersection_id_set)} intersection item_id, remove them in negative")
            for item_id in intersection_id_set:
                negative_dataset.pop(item_id)
        assert len(positive_dataset) > 0 and len(negative_dataset) > 0, "positive and negative must not be empty"
        self.logger.info(
            f"positive dataset size: {len(positive_dataset)}, negative dataset size: {len(negative_dataset)}")
        # 合并正负例，生成训练feature
        positive_features = np.array(list(positive_dataset.values()))
        negative_features = np.array(list(negative_dataset.values()))
        features = np.concatenate([positive_features, negative_features], axis=0)
        # 生成训练label
        labels = np.concatenate([np.ones(len(positive_features)), np.zeros(len(negative_features))], axis=0)
        self.model = self.model.fit(features, labels)
        self.logger.info(f"train {self.model_name} successfully")

    def infer(self, datasets):
        self.logger.info(f"start infer {self.model_name}")
        infer_result = {}
        for dataset_key, dataset in datasets.items():
            features = np.array(list(dataset.values()))
            scores = self.model.predict_proba(features)[:, -1]
            scores = scores.flatten().tolist()
            infer_result[dataset_key] = dict(zip(dataset.keys(), scores))
        self.logger.info(f"infer {self.model_name} finished")
        return infer_result
